import threading
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler
from typing import Callable, Mapping

from common.types import Lock
from infection_monkey.exploit import IAgentBinaryRepository, RetrievalError

from .agent_binary_request import AgentBinaryDownloadReservation, ReservationID


class AgentBinaryHTTPRequestHandler(BaseHTTPRequestHandler):
    @classmethod
    def reserve_download(cls, reservation: AgentBinaryDownloadReservation):
        raise NotImplementedError()

    @classmethod
    def clear_reservation(cls, reservation_id: ReservationID):
        raise NotImplementedError()


class ThreadingHTTPHandlerFactory:
    def __init__(self, agent_binary_repository: IAgentBinaryRepository):
        self.agent_binary_repository = agent_binary_repository

    def __call__(self):
        return get_http_handler(self.agent_binary_repository, {}, {}, lambda: threading.Lock())


def get_http_handler(
    agent_binary_repository: IAgentBinaryRepository,
    reservations: Mapping[ReservationID, AgentBinaryDownloadReservation],
    locks: Mapping[ReservationID, Lock],
    create_lock: Callable[[], Lock],
):
    def reserve_download(cls, reservation: AgentBinaryDownloadReservation):
        if reservation.id in cls.reservations:
            raise KeyError(f"Request ID {reservation.id} is already registered")
        cls.reservations[reservation.id] = reservation
        cls.locks[reservation.id] = create_lock()

    def clear_reservation(cls, reservation_id: ReservationID):
        del cls.reservations[reservation_id]

    return type(
        "AgentBinaryHTTPHandler",
        (AgentBinaryHTTPRequestHandler,),
        {
            "agent_binary_repository": agent_binary_repository,
            "reservations": reservations,
            "locks": locks,
            "do_GET": _do_GET,
            "reserve_download": classmethod(reserve_download),
            "clear_reservation": classmethod(clear_reservation),
        },
    )


def _do_GET(self):
    cls = self.__class__
    reservation_id = ReservationID(self.path.split("/")[-1])  # Parse request from the URL

    try:
        lock = cls.locks[reservation_id]
    except KeyError:
        self.send_response(404)
        self.end_headers()
        raise

    with lock:
        reservation: AgentBinaryDownloadReservation = cls.reservations[reservation_id]
        if reservation.download_completed.is_set():
            self.send_error(
                HTTPStatus.TOO_MANY_REQUESTS,
                "A download has already been requested",
            )

        try:
            agent_binary = cls.agent_binary_repository.get_agent_binary(
                reservation.operating_system
            )
        except RetrievalError:
            self.send_error(
                HTTPStatus.INTERNAL_SERVER_ERROR, "The binary does not exist on the server"
            )

        self.send_response(HTTPStatus.OK)
        self.send_header("Content-type", "application/octet-stream")
        self.end_headers()

        bytes_to_send = reservation.transform_agent_binary(agent_binary.getvalue())

        self.wfile.write(bytes_to_send)
        reservation.download_completed.set()
